# Copyright (c) 2022 Presto Labs Pte. Ltd.
# Author: qiangluo

import collections
import hashlib
import hmac
import json
import os
import requests.auth
import time
import urllib.parse

import pyotp
from absl import app, flags


def _get_if_not(d, y):
  if y in d:
    return d[y]
  else:
    return None


class AuthKey(object):
  @staticmethod
  def from_file(key_file):
    assert os.path.exists(key_file), key_file

    key = json.load(open(key_file))
    owner = None
    if owner is None:
      owner = _get_if_not(key, "owner")
    if owner is None:
      owner = _get_if_not(key, "account_name")

    access_key = None
    if access_key is None:
      access_key = _get_if_not(key, "access_key")
    if access_key is None:
      access_key = _get_if_not(key, "api_key")
    if access_key is None:
      access_key = _get_if_not(key, "api-key")

    secret_key = None
    if secret_key is None:
      secret_key = _get_if_not(key, "secret_key")
    if secret_key is None:
      secret_key = _get_if_not(key, "api_secret")
    if secret_key is None:
      secret_key = _get_if_not(key, "api-secret")

    return AuthKey(
        owner=owner,
        key_name=_get_if_not(key, "name"),
        access_key=access_key,
        secret_key=secret_key,
        generated=_get_if_not(key, "generated"),
        generated_human=_get_if_not(key, "generated_human"),
        raw=key,
        key_file=key_file,
    )

  def __init__(self,
               *,
               owner=None,
               key_name=None,
               access_key=None,
               secret_key=None,
               generated=None,
               generated_human=None,
               raw=None,
               key_file=None):
    self._owner = owner
    self._key_name = key_name
    assert type(access_key) == str, access_key
    assert type(secret_key) == str, secret_key
    self._access_key = access_key
    self._secret_key = secret_key
    self._secret_key_bytes = secret_key.encode("utf-8")
    self._secret_key_upper_bytes = secret_key.upper().encode("utf-8")
    self._generated = generated
    self._generated_human = generated_human
    self._raw = raw
    self._key_file = key_file

  @property
  def owner(self):
    return self._owner

  @property
  def access_key(self):
    return self._access_key

  @property
  def secret_key(self):
    return self._secret_key

  @property
  def secret_key_bytes(self):
    return self._secret_key_bytes

  @property
  def secret_key_upper_bytes(self):
    return self._secret_key_upper_bytes

  @property
  def api_key(self):
    return self._access_key

  @property
  def api_secret(self):
    return self._secret_key

  @property
  def key_file(self):
    return self._key_file

  @property
  def refresh_token(self):
    return self._raw.get("refresh_token", None)

  def get_value(self, key):
    # do not add default value
    return self._raw[key]

  def as_json(self):
    o = collections.OrderedDict()
    o["owner"] = self._owner
    o["name"] = self._key_name
    o["access_key"] = self._access_key
    o["secret_key"] = self._secret_key
    if "refresh_token" in self._raw:
      o["refresh_token"] = self._raw["refresh_token"]
    o["generated"] = self._generated
    o["generated_human"] = self._generated_human

    return o

  def as_json_str(self):
    return json.dumps(self.as_json(), indent=2)

  def save_file(self, key_file):
    val = self.as_json_str()
    with open(key_file, "w") as f:
      print(val, file=f)


def is_binance_stream_url(url):
  return '/api/v1/userDataStream' in url or '/fapi/v1/listenKey' in url


class BinanceAuth(requests.auth.AuthBase):
  def __init__(self, key_file):
    self._key = AuthKey.from_file(key_file)

  def _gen_signature(self, data: bytes):
    signature = hmac.new(self._key.secret_key_bytes, data, digestmod=hashlib.sha256)
    return signature.hexdigest()

  def __call__(self, req: requests.PreparedRequest):
    if not is_binance_stream_url(req.url):
      query = urllib.parse.urlparse(req.url).query
      signature = self._gen_signature(query.encode())
      req.url = '%s&signature=%s' % (req.url, signature)
    req.headers.update({'X-MBX-APIKEY': self._key.access_key})
    return req


class FTXAuth(requests.auth.AuthBase):
  def __init__(self, key_file):
    self._key = AuthKey.from_file(key_file)

  def __call__(self, req: requests.PreparedRequest):
    # assert req.method == 'GET', req.method
    timestamp = int(time.time() * 1000)
    signature_payload = f'{timestamp}{req.method}{req.path_url}'.encode()
    if req.body:
      signature_payload += req.body
    signature = hmac.new(self._key.api_secret.encode(), signature_payload, 'sha256').hexdigest()
    req.headers.update({
        'FTX-KEY': self._key.api_key,
        'FTX-TS': str(timestamp),
        'FTX-SIGN': signature,
    })
    if self._key._raw.get('account_type') != 'master':
      req.headers.update({'FTX-SUBACCOUNT': self._key.owner})
    return req


def generate_api_key_signature(secret, verb, url, expires, data):
  parsed_url = urllib.parse.urlparse(url)
  path = parsed_url.path
  if parsed_url.query:
      path = path + '?' + parsed_url.query
  if isinstance(data, (bytes, bytearray)):
      data = data.decode('utf8')
  message = verb + path + str(expires) + data
  signature = hmac.new(
      bytes(secret, 'utf8'), bytes(message, 'utf8'), digestmod=hashlib.sha256)
  return signature.hexdigest()


class PrexAuth(requests.auth.AuthBase):
  def __init__(self, key_file):
    self._key = AuthKey.from_file(key_file)

  def _gen_signature(self, data):
    signature = hmac.new(self._key.secret_key_bytes, data, digestmod=hashlib.sha256)
    return signature.hexidgest()

  def __call__(self, req):
    expires = int(round(time.time()) + 5)
    req.headers['api-expires'] = str(expires)
    req.headers['api-key'] = self._key.access_key
    req.headers['api-signature'] = generate_api_key_signature(
        self._key.secret_key, req.method, req.url, expires, req.body or '')
    return req


class SyncQueryer(object):
  def __init__(self, *, api_host, auth, timeout=10):
    self._timeout = timeout
    self._auth = auth
    self._api_host = api_host

  def query(self, *, method, path, params=None, data=None):
    assert method in ('GET', 'PUT', 'POST', 'DELETE'), method
    url = urllib.parse.urljoin(self._api_host, path)
    response = None
    try:
      response = requests.request(
          method,
          url,
          params=params,
          auth=self._auth,
          timeout=self._timeout,
          data=data)
      content = response.json()
      return content
    except requests.exceptions.RequestException as e:
      print('RequestException: %s', e)
      raise
    except json.JSONDecodeError as e:
      print('Json decode: %s\n %s', response.content, e)
      print(response)
      print(response.text)
      raise
    except Exception as e:
      print('Unknown exception, %s', e)
      raise


class BinancePrivateClient():
    def __init__(self, key_file):
        self.key_file = key_file
        self.url = 'https://api.binance.com'
        self.auth = BinanceAuth(key_file)
        self.query = SyncQueryer(api_host=self.url, auth=self.auth)

    def withdraw(self, *, coin, address, amount):
        method = "POST"
        path = "/sapi/v1/capital/withdraw/apply"
        params = {
            'timestamp': int(time.time() * 1000),
            "coin": coin,
            "address": address,
            "amount": amount
        }
        print(params)
        msg = self.query.query(method=method, path=path, params=params)
        return msg

    def deal_history(self):
        method = "GET"
        path = "/sapi/v1/capital/withdraw/history"
        params = {
            'timestamp': int(time.time() * 1000),
        }
        msg = self.query.query(method=method, path=path, params=params)
        return msg


class FtxPrivateClient():
    def __init__(self, key_file):
        self.key_file = key_file
        self.url = 'https://ftx.com'
        self.auth = FTXAuth(key_file)
        self.query = SyncQueryer(api_host=self.url, auth=self.auth)

    # The exchange does not
    def deal_history(self):
        method = "GET"
        path = "/api/wallet/withdrawals"
        params = {}
        msg = self.query.query(method=method, path=path, params=params)
        return msg

    def withdraw(self, *, coin, address, amount):
        method = "POST"
        path = "/api/wallet/withdrawals"
        params = {
            "coin": coin,
            "address": address,
            "size": amount
        }
        msg = self.query.query(method=method, path=path, params=params)
        return msg


class PrexPrivateClient():
    def __init__(self, key_file):
      self.key_file = key_file
      self.url = 'https://api.sandbox.prex.run'
      self.auth = PrexAuth(key_file)
      self.query = SyncQueryer(api_host=self.url, auth=self.auth)

    def deal_history(self, count=10):
      params = {"count": count}
      method = "GET"
      path = "/api/v1/wallet/history"
      msg = self.query.query(method=method, path=path, params=params)
      return msg

    def withdraw(self, *, coin, address, amount):
      method = "POST"
      key = "A5QN4EKNVTJDYPNA"
      totp = pyotp.TOTP(key)
      path = "/api/v1/wallet/withdrawal"
      body = {
          "address": address,
          "addressParams": "{}",
          "blockchainName": "bitcoin",
          "currency": coin,
          "fee": "0.0005",  # 这里是要先根据币的种类查询fee吗？
          "netAmount": amount,
          "token": str(totp.now()) # 不确定
      }
      msg = self.query.query(method=method, path=path, data=body)
      return msg


def main(_):
  exchange = flags.FLAGS.exchange
  keypath = flags.FLAGS.keypath
  coin = flags.FLAGS.coin
  address = flags.FLAGS.address
  amount = flags.FLAGS.amount
  history = flags.FLAGS.history

  try:
    assert keypath is not None, 'keypath not set'
    assert exchange in ('Ftx', 'Binance', 'Prex', 'Aqx'), exchange
    assert all([coin, address, amount]) or history, [coin, address, amount, history]
  except:
    print_help()
    raise

  if exchange == "Ftx":
    client = FtxPrivateClient(keypath)
  elif exchange == "Binance":
    client = BinancePrivateClient(keypath)
  elif exchange in ("Prex", 'Aqx'):
    client = PrexPrivateClient(keypath)
  else:
    raise Exception("Unknown exchange: " + exchange)

  if history:
    msg = client.deal_history()
    print(json.dumps(msg, indent=2))
  else:
    msg = client.withdraw(coin=coin, address=address, amount=amount)
    print(json.dumps(msg, indent=2))


def print_help():
  help_msg = """
example
query historical transaction records
python ./withdraw_tool.py --exchange=Binance --keypath=path/to/key --history

withdraw
python ./withdraw_tool.py --exchange=Binance --keypath=path/to/key --coin=BTC --address=eeffaa09xxyyzz --amount=100
"""
  print(help_msg)


if __name__ == '__main__':
  flags.DEFINE_string(name='exchange', help='exchange type', default=None)
  flags.DEFINE_string(name='keypath', help='Path to key file', default=None)
  flags.DEFINE_string(name='coin', help='Coin type, such as BTC,ETH,USDT', default=None)
  flags.DEFINE_string(name='address', help='Withdrawals address', default=None)
  flags.DEFINE_integer(name='amount', help='The amount of withdrawals', default=None)
  flags.DEFINE_bool(name="history", help="Query historical transaction records", default=None)
  app.run(main)
